import base64
from functools import wraps
import re
import subprocess
import os
from datetime import datetime
import sys
import time
import traceback
from Config import adb
from termcolor import colored
import os
import subprocess
import xml.etree.ElementTree as ET
from pydantic import BaseModel, Field, validator
from typing import Literal, Union, Optional
from PIL import Image, ImageDraw
from PIL import ImageFont
import concurrent.futures
import cv2
import numpy as np
# ------
class ADBShell:
    def __init__(self):
        # Launch the ADB shell once
        self.init_proc()
    def init_proc(self):
        self.proc = subprocess.Popen(
            [adb, 'shell'],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True)
    def send_command(self, cmd: str, **kwargs) -> dict:
        """
        Write a command to the persistent ADB shell and read the output.
        Use a sentinel marker so we know when to stop reading.
        """
        # sentinel = "___SENTINEL___"
        print("start send command")
        if self.proc.poll() is not None:
            self.init_proc()
        print("start write command")
        # try:
        self.proc.stdin.write(cmd + '\n')
        # self.proc.stdin.write(f"echo {sentinel}\n")
        self.proc.stdin.flush()
        # except OSError as e:
        #     if e.errno == 22:  # Invalid argument
        #         self.init_proc()
        #         self.proc.stdin.write(cmd + '\n')
        #         self.proc.stdin.write(f"echo {sentinel}\n")
        #         self.proc.stdin.flush()
        #         return {'returncode':0,'stdout':"",'stderr':""}
        #     else:
        #         raise e
        # line = self.proc.stdout.readline()
        # print(line)
        # errline = self.proc.stderr.readline()
        # print(errline)
        # print("end")
        return {'returncode':0,'stdout':"",'stderr':""}
    def close(self):
        if self.proc:
            self.proc.terminate()

adb_shell = ADBShell()
#  adb install ADBKeyBoard.apk
#   adb shell ime enable com.android.adbkeyboard/.adbIme
#   adb shell ime set com.android.adbkeyboard/.adbIme
def timeit(color='magenta'):
    """装饰器：测量函数执行时间，并允许指定输出颜色"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            print(colored(f"{func.__name__} time:{(end_time - start_time):.3f}s", color))
            sys.stdout.flush()
            return result
        return wrapper
    return decorator
# Cache for resolution values
_resolution_cache = None
# @timeit('red')
def subrun(cmd: str, **kwargs):
    """
    Send a command to the persistent ADB shell, then gather
    the output, error, and return code in a structure that
    'getresult' can handle.
    """
    return subprocess.run(f"{adb} shell {cmd}",**kwargs)
    # return adb_shell.send_command(cmd, **kwargs)
    # out = adb_shell.send_command(cmd, **kwargs)
    # code_str = adb_shell.send_command("echo $?")
    # try:
    #     return_code = int(code_str.strip() or "1")
    # except ValueError:
    #     return_code = 1 
    # return {'returncode':return_code,'stdout':out,'stderr':"" if return_code==0 else "Unknown error or nonzero exit code"}
class TaskComplete(BaseModel):
    """
    no params.if you see the task is complete,call this tool to terminate the task
    """
    pass
class checkComplete(BaseModel):
    """
    no params.If you want to check if the task is complete,call this tool, then user will send you a screenshot and a11yTree information.
    """
    pass
class ClickParams(BaseModel):
    x: int = Field(..., description="X coordinate for tap")
    y: int = Field(..., description="Y coordinate for tap")
    tell_user_speech:str=Field(...,description="let the textToSpeech system tell user what you are doing.Use Chinese,no more than 10 words")

class LongClickParams(BaseModel):
    x: int = Field(..., description="X coordinate for long press")
    y: int = Field(..., description="Y coordinate for long press")
    duration: int = Field(..., description="Duration of long press in milliseconds")
    tell_user_speech:str=Field(...,description="let the textToSpeech system tell user what you are doing.Use Chinese,no more than 10 words")

class SwipeParams(BaseModel):
    x: int = Field(..., description="Starting X coordinate")
    y: int = Field(..., description="Starting Y coordinate")
    dir: Literal["up", "down", "left", "right"] = Field(..., description="Direction of swipe")
    dist: Literal["long", "mid", "short"] = Field(..., 
        description="Distance of swipe (long=0.7, mid=0.4, short=0.15 of screen)")
    speed: Literal["vfast", "fast", "normal", "slow"] = Field(..., 
        description="Speed of swipe (vfast=0.5, fast=1, normal=1.5, slow=2.5)")
    tell_user_speech:str=Field(...,description="let the textToSpeech system tell user what you are doing.Use Chinese,no more than 10 words")
    
class TextParams(BaseModel):
    text: str = Field(..., description="Text to input.Note that before calling TextParams, you must call ClickParams to click the inputArea so that it gets focus")
    tell_user_speech:str=Field(...,description="let the textToSpeech system tell user what you are doing.Use Chinese,no more than 10 words")

class KeycodeParams(BaseModel):
    code: int = Field(..., description="Android keycode to send")
    tell_user_speech:str=Field(...,description="let the textToSpeech system tell user what you are doing.Use Chinese,no more than 10 words")
class ScreenshotParams(BaseModel):
    """get the screenshot and a11y tree of the current window (no parameters required)"""
    pass
    # ref_h_pix:Optional[int] = Field(200, description="Draw horizontal reference lines to screenshot. value is pixel interval of horizontal lines")
    # ref_v_pix:Optional[int] = Field(200, description="Draw vertical reference lines to screenshot. value is pixel interval of vertical lines")
# class a11yTreeParams(BaseModel):
#     """get the accessibility tree of the current window (no parameters required)"""
#     pass
# class speakParams(BaseModel):
#     text:str=Field(...,description="FOR each conversation round (unless the task is complete, identified by you not calling any tool), you must call this tool to let voice generation system tell user what you are doing.The text should be in Chinese, no more than 10 words!!")
#     pass
AGENT_TOOLS=[ClickParams,LongClickParams,SwipeParams,TextParams]

def getresult(result,addResult:dict|None=None):
    if addResult is None:
        addResult={'success':True,'content':""}
    if result.returncode == 0:
        return {'success':addResult.get('success'),'content':f"{addResult.get('content')} {result.stdout}"}
    else:
        return {'success':False,'content':f"{addResult.get('content')} {result.stderr}"}
def router(tool,params):
    try:   
        if tool=="ClickParams":
            return click(params.get('x'),params.get('y'))
        elif tool=="LongClickParams":
            return longclick(params.get('x'),params.get('y'),params.get('duration'))
        elif tool=="SwipeParams":
            return swipe(int(params.get('x')),int(params.get('y')),int(params.get('dir')),int(params.get('dist')),int(params.get('speed')))
        elif tool=="TextParams":
            # if all(ord(c) < 128 for c in params.get('text')):
            return text(params.get('text'))
            # else:
            #     return {'success': False, 'content': 'Error: can only text English/ASCII characters'}
        elif tool=="KeycodeParams":
            return keycode(params.get('code'))
        elif tool=="ScreenshotParams":
            return screenAndXML()
        elif tool=="checkComplete":
            return {'success':True,'content':"please view the info sent by user and tell me if the task is complete"}
    except Exception as e:
        return {'success': False, 'content': str(e)+traceback.format_exc()}
def loc2XY(loc):
    left, top, right, bottom = map(int,loc.strip('[]').replace('][',',').split(','))
    X = (left + right) // 2
    Y = (top + bottom) // 2
    return [X,Y]
def resolution():
    global _resolution_cache
    if _resolution_cache is not None:
        return _resolution_cache
    result = getresult(subprocess.run(f'{adb} shell wm size',stdout=subprocess.PIPE,stderr=subprocess.PIPE,text=True))
    if result.get('success'):
        resolution = result.get('content').strip().split(' ')[-1]  # Extract resolution
        width, height = resolution.split('x')  # Split width and height
        _resolution_cache = (float(width), float(height))
        return {'success':True,'content':_resolution_cache}
    else:
        return result
def click(x,y):
    cmd = f'input tap {x} {y}'
    return getresult(subrun(cmd))
def longclick(x, y, duration):
    cmd = f"input touchscreen swipe {x} {y} {x} {y} {duration}"
    return getresult(subrun(cmd))
def swipe(x,y,dir,dist,speed):
    width,height=resolution()
    distmap = {'long': 0.7, 'mid': 0.4, 'short': 0.15}
    ndist=distmap.get(dist)
    speedmap= {'vfast':0.5 ,'fast':1.0,'normal':1.5,'slow':2.5}
    nspeed=speedmap.get(speed)
    duration=1000*ndist*nspeed
    xdist=width*ndist
    ydist=height*ndist
    if dir=='up':
        y2=y-ydist
        x2=x
    elif dir=='down':
        y2=y+ydist
        x2=x
    elif dir=='left':
        x2=x-xdist
        y2=y
    elif dir=='right':
        x2=x+xdist
        y2=y
    if x2>width:
        x2=width
    if x2<0:
        x2=0
    if y2>height:
        y2=height
    if y2<0:
        y2=0
    cmd = f"input touchscreen swipe {x} {y} {int(x2)} {int(y2)} {int(duration)}"
    return getresult(subrun(cmd))
def text(input_text: str):
    """
    Intelligent 'text' function that either uses:
    - standard `adb shell input text` for basic ASCII
    - or ADBKeyBoard broadcast for non-ASCII (Chinese, etc.)
    Then reverts input method to default at the end.
    """
    if all(ord(c) < 128 for c in input_text):
        # ASCII path (original)
        text_value = input_text.replace(" ", "%s")  # preserve spacing
        res = getresult(subrun(f"input text '{text_value}'"))
        return res
    else:
        # Non-ASCII path
        time.sleep(0.5)
        res=getresult(subrun(f"am broadcast -a ADB_INPUT_TEXT --es msg '{input_text}'"))
        return res
def keycode(code):
    return getresult(subrun(f'input keyevent {str(code)}'))
@timeit('blue')
def screenshotRaw(getPath=True)->str|dict|None:
    photos_dir = os.path.join(os.getcwd(), 'PHOTOS')
    if not os.path.exists(photos_dir):
        os.makedirs(photos_dir)
    localPath = os.path.join(photos_dir, f'screenshot.png')
    tempPath = '/sdcard/screenshot.png'
    result=getresult(subrun(f"screencap -p {tempPath}"))
    result=getresult(subprocess.run(f"{adb} pull {tempPath} {localPath}"),result)
    if result.get('success'):
        if getPath:
            return localPath
        return {'success':True,'content':localPath}
    if getPath:
        return None
    return result
@timeit('blue')
def draw_reference_lines(localPath: str, n_width: int, m_height: int) -> Image.Image:
    """
    1) Load the original image.
    2) Resize it based on a fixed new_height (keeping aspect ratio).
    3) For each original coordinate multiple (x, y) in n_width / m_height steps,
       compute scaled_x, scaled_y, and draw lines/labels in the resized image.
    4) Save and return the resized image with reference lines.
    """

    # Read the original image
    img = cv2.imread(localPath)
    orig_height, orig_width = img.shape[:2]

    # --- Resize first ---
    new_height = 900
    aspect_ratio = orig_width / orig_height
    new_width = int(new_height * aspect_ratio)

    # Resize the image to new_width x new_height
    img = cv2.resize(img, (new_width, new_height))

    # Compute scaling factors so we can map original coords to the resized coords
    scale_x = new_width / orig_width
    scale_y = new_height / orig_height

    # Prepare drawing parameters
    color = (0, 0, 255)  # Red in BGR
    thickness = 1
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.8*scale_x

    # --- Draw lines and coordinate labels in the resized image ---
    if m_height:
        # Horizontal lines
        for y in range(0, orig_height, m_height):
            scaled_y = int(y * scale_y)
            cv2.line(img, (0, scaled_y), (new_width, scaled_y), color, thickness)

    if n_width:
        # Vertical lines
        for x in range(0, orig_width, n_width):
            scaled_x = int(x * scale_x)
            cv2.line(img, (scaled_x, 0), (scaled_x, new_height), color, thickness)

    # Intersection labels (original coords, placed at scaled positions)
    if m_height and n_width:
        for y in range(0, orig_height, m_height):
            for x in range(0, orig_width, n_width):
                scaled_x = int(x * scale_x)
                scaled_y = int(y * scale_y)
                coord_text = f"{x},{y}"
                text_pos = (scaled_x + 1, max(scaled_y - 1, 0))
                cv2.putText(img, coord_text, text_pos, font, font_scale, color, thickness)

    # Write the final image (overwriting the original file, if desired)
    cv2.imwrite(localPath, img)

    # Convert to PIL Image (if you need to return or reuse)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(img_rgb)
    return pil_image
@timeit('blue')
def screenshotDrawline(getPath:bool,n_width:int, m_height:int,addLine:bool)->str|dict:
    """
    Draw reference grid lines on the screenshot.
    
    Args:
        image_path: Path to the screenshot image
        n_width: Draw vertical line every N pixels in width
        m_height: Draw horizontal line every M pixels in height
    """
    try:
        pathDict = screenshotRaw(getPath=False)
        assert isinstance(pathDict,dict)
        if pathDict.get('success'):
            localPath:str=pathDict.get('content')#type:ignore
            if addLine:
                draw_reference_lines(localPath,n_width,m_height)
            if getPath:
                return localPath
            return {'success':True,'content':localPath}
        if getPath:
            return ""
        return pathDict
    except Exception as e:
        if getPath:
            return ""
        return {'success':False,'content':str(e)+traceback.format_exc()}
@timeit('green')
def xmlshot(getPath: bool) -> str | dict:
    dir = 'scriptXML'
    os.makedirs(dir, exist_ok=True)
    name = 'scriptxml.xml'
    path = os.path.join(dir, name)

    # Get UI dump directly with exec-out to avoid file operations
    cmd = [adb, 'exec-out', 'uiautomator', 'dump', '/dev/tty']
    start_time = time.time()
    result = subprocess.run(
        cmd,
        check=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=False,  # Changed to handle binary data
    )
    end_time = time.time()
    print(f"xmldump time:{(end_time - start_time):.3f}s")
    
    if result.returncode == 0 and result.stdout:
        try:
            # Try to decode as UTF-8
            xml_content = result.stdout.decode('utf-8').strip()
            # Write directly to local file
            with open(path, 'w', encoding='utf-8') as f:
                f.write(xml_content.replace("UI hierchary dumped to: /dev/tty",""))
            if getPath:
                return path
            return {'success': True, 'content': path}
        except UnicodeDecodeError:
            # Fallback if we can't decode the output
            return {'success': False, 'content': "Error: Unable to decode XML output"}
    
    if getPath:
        return ""
    return {'success': False, 'content': result.stderr.decode('utf-8', errors='replace') if result.stderr else "No output from uiautomator"}
@timeit('green')
def traverse(xmlpath):
    try:
        tree = ET.parse(xmlpath)    
    except Exception as e:
        with open(xmlpath,'r',encoding='utf-8') as f:
            xml_content = f.read()
        print(xml_content)
        print(e)
        return [],f"error getting xml {e}"
    root = tree.getroot()
    index=0
    info = []
    infostr=[]
    for elem in root.iter():  
        text = str(elem.get('text'))
        desc=str(elem.get('content-desc')) 
        if text in desc:
            content=desc
        elif desc in text:
            content=text
        else:
            content=text+','+desc
        classes=str(elem.get('class')).replace('android','').replace('widget','').replace('View','').replace('view','').replace('app','').replace('.','')
        property=[key for key, value in elem.attrib.items() if value == 'true' and key!='enabled' and key!='focusable']            
        # Map property names to their shortened versions
        property_mapping = {
            'scrollable': 'scl',
            'clickable': 'clk',
            'long-clickable': 'lclk',
            'editable': 'edit',
            'password': 'pwd',
            'checkable': 'check',
            'checked': 'checked'
        }
        # Filter out 'focusable' and transform others using the mapping
        property = [property_mapping.get(prop, prop) for prop in property if prop != 'focusable' and prop != 'enabled']
        if (property!=[] or content!=''):
            bounds = elem.get('bounds')
            if bounds and bounds!='[0,0][0,0]':
                info.append((index,content,classes,property,(loc2XY(bounds)[0],loc2XY(bounds)[1])))
                args=[f"{index}:"]
                if content!='':
                    args.append(f"{content},")
                if classes!='':
                    args.append(f"{classes},")
                if property!=[]:
                    args.append(f"{'|'.join(property)},")
                args.append(f'({loc2XY(bounds)[0]},{loc2XY(bounds)[1]})')
                infostr.append("".join(args))
                # print("".join(args))
                index=index+1 
    return info,"\n".join(infostr)
def formatlist(infolist:list[tuple[int,str,str,list[str],tuple[int,int]]],fmtStr:str)->list[str]:
    fmtlist:list[str]=[]
    for item in infolist:
        i, con, cl, pro, loc = item 
        fmtitem = fmtStr.format(i=i, con=con, cl=cl, pro=pro, loc=loc)
        fmtlist.append(fmtitem)
    return fmtlist 
def findActivity():
    result = subprocess.run([adb, 'shell', 'dumpsys', 'activity', 'activities'], stdout=subprocess.PIPE)
    output = result.stdout.decode()
    pattern = re.compile(r".*mFocusedActivity|ResumedActivity|FocusedActivity|topResumedActivity=ActivityRecord{.*?(com.*?)}")
    matches = pattern.findall(output)
    print(f'current main activity is {matches[0]}')
    ans=matches[0].replace('.','_').replace('/','_')
    return ans
@timeit('green')
def xmlshotTraverse(getStr:bool)->str|dict:
    """
    Helper that calls xmlshot and then traverse, returning { 'info': [...], 'a11y': '...' } 
    (or you can return whatever data structure you prefer).
    """
    xml_path_dict = xmlshot(getPath=False)  # returns the path to the XML
    assert isinstance(xml_path_dict,dict)
    if xml_path_dict.get('success'):
        xml_path:str=xml_path_dict.get('content')#type:ignore
        info_list, infostr = traverse(xml_path)
        if getStr:
            return infostr
        return {'success':True,'content':infostr}
    else:
        if getStr:
            return ""
        return xml_path_dict
@timeit('yellow')
def screenAndXML()->dict:
    """
    New method to:
      1) Use screenshotDrawline(...) to capture an image (with lines automatically drawn if you want).
      2) In parallel, run xmlshot+traverse.
      3) Return both results in a simple dictionary without nested if-checks.
    """
    ref_h_pix=200
    ref_v_pix=200
    with concurrent.futures.ThreadPoolExecutor() as executor:
        work1 = executor.submit(screenshotDrawline,False, ref_h_pix, ref_v_pix,True)
        work2 = executor.submit(xmlshotTraverse,False)
        pathDict=work1.result()
        infostrDict=work2.result()
        assert isinstance(pathDict,dict)
        assert isinstance(infostrDict,dict)
    return {
        'success':pathDict.get('success') and infostrDict.get('success'),
        'path': pathDict.get('content',""),
        'a11y': infostrDict.get('content',"")
    }
def adbIme():
    return getresult(subrun("ime set com.android.adbkeyboard/.AdbIME"))
def defaultIme():
    return getresult(subrun("ime set com.vivo.ai.ime.nex/com.vivo.ai.ime.main.IMEService"))
if __name__ == '__main__':
    # infolist,infostr=traverse(xmlshot())
    # print(str(screenAndXML()).replace('\\n', '\n'))
    adbIme()
    text("AI：现在支持中文输入了")
    defaultIme()
adb_shell.close()






